#!/usr/bin/env python
# coding: utf-8

# In[5]:


import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import MinMaxScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.metrics import roc_curve, auc
import os

# 设置 Tushare token
ts.set_token('7705d187ddde476492fc7c72b674218f8d383448f7126a0457392231')
pro = ts.pro_api()

# 1. 获取两年至少10支股票的日线行情数据
def get_stock_data():
    # 选取 10 支股票
    stock_list = ['000001.SZ', '000002.SZ', '000004.SZ', '000005.SZ', '000006.SZ',
                  '000007.SZ', '000008.SZ', '000009.SZ', '000010.SZ', '000011.SZ']
    start_date = '20230508'
    end_date = '20250508'
    all_data = []
    for stock in stock_list:
        df = pro.daily(ts_code=stock, start_date=start_date, end_date=end_date)
        all_data.append(df)
    data = pd.concat(all_data, ignore_index=True)
    return data

# 2. 对股票行情数据进行分类（打标签）与技术指标计算
def add_features_and_labels(data):
    # 打标签：如果次日收盘价上涨，标签为 1，否则为 0
    data['label'] = (data.groupby('ts_code')['close'].shift(-1) > data['close']).astype(int)
    # 计算技术指标
    data['ma5'] = data.groupby('ts_code')['close'].rolling(window=5).mean().reset_index(0, drop=True)
    data['ma10'] = data.groupby('ts_code')['close'].rolling(window=10).mean().reset_index(0, drop=True)
    data['ma20'] = data.groupby('ts_code')['close'].rolling(window=20).mean().reset_index(0, drop=True)
    data['std5'] = data.groupby('ts_code')['close'].rolling(window=5).std().reset_index(0, drop=True)
    data['std10'] = data.groupby('ts_code')['close'].rolling(window=10).std().reset_index(0, drop=True)
    data['std20'] = data.groupby('ts_code')['close'].rolling(window=20).std().reset_index(0, drop=True)
    data['pct_change'] = data.groupby('ts_code')['close'].pct_change()
    data['volume_pct_change'] = data.groupby('ts_code')['vol'].pct_change()
    data['high_low_ratio'] = data['high'] / data['low']
    data['close_open_ratio'] = data['close'] / data['open']
    data['close_ma5_ratio'] = data['close'] / data['ma5']
    data['close_ma10_ratio'] = data['close'] / data['ma10']
    data['close_ma20_ratio'] = data['close'] / data['ma20']
    data['rsi'] = calculate_rsi(data).reset_index(0, drop=True)
    macd, signal = calculate_macd(data)
    data['macd'] = macd.reset_index(0, drop=True)
    data['signal'] = signal.reset_index(0, drop=True)
    boll_upper, boll_middle, boll_lower = calculate_bollinger_bands(data)
    data['boll_upper'] = boll_upper.reset_index(0, drop=True)
    data['boll_middle'] = boll_middle.reset_index(0, drop=True)
    data['boll_lower'] = boll_lower.reset_index(0, drop=True)
    return data

def calculate_rsi(data, period=14):
    delta = data.groupby('ts_code')['close'].diff()
    up = delta.clip(lower=0)
    down = -delta.clip(upper=0)
    avg_gain = up.rolling(window=period).mean()
    avg_loss = down.rolling(window=period).mean()
    rs = avg_gain / avg_loss
    rsi = 100 - (100 / (1 + rs))
    return rsi

def calculate_macd(data, short_window=12, long_window=26, signal_window=9):
    short_ema = data.groupby('ts_code')['close'].ewm(span=short_window, adjust=False).mean()
    long_ema = data.groupby('ts_code')['close'].ewm(span=long_window, adjust=False).mean()
    macd = short_ema - long_ema
    signal = macd.ewm(span=signal_window, adjust=False).mean()
    return macd, signal

def calculate_bollinger_bands(data, window=20, std_dev=2):
    middle_band = data.groupby('ts_code')['close'].rolling(window=window).mean()
    std = data.groupby('ts_code')['close'].rolling(window=window).std()
    upper_band = middle_band + (std * std_dev)
    lower_band = middle_band - (std * std_dev)
    return upper_band, middle_band, lower_band

# 3. 对分类计算后的数据进行建模前的处理与分析
def preprocess_data(data):
    # 处理空值
    data = data.dropna()
    # 选择特征和标签
    features = data.drop(['ts_code', 'trade_date', 'label'], axis=1)
    labels = data['label']
    # 归一化
    scaler = MinMaxScaler()
    features_scaled = scaler.fit_transform(features)
    # 主成分分析
    pca = PCA(n_components=0.95)
    features_pca = pca.fit_transform(features_scaled)
    # 相关性分析
    corr_matrix = features.corr()
    plt.figure(figsize=(12, 8))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm')
    plt.title('Correlation Matrix')
    plt.savefig('correlation_matrix.png')
    plt.show()
    # 降维后的特征
    features_pca_df = pd.DataFrame(features_pca, columns=[f'PC{i+1}' for i in range(features_pca.shape[1])])
    # 个股画像（选择 6 个属性）
    stock_profile = data[['ts_code', 'close', 'ma5', 'ma10', 'ma20', 'volume_pct_change', 'rsi']]
    
    # 数据均衡 - 修改为可选SMOTE或简单跳过
    try:
        from imblearn.over_sampling import SMOTE
        smote = SMOTE()
        features_resampled, labels_resampled = smote.fit_resample(features_pca, labels)
        print("使用SMOTE进行数据均衡")
    except ImportError:
        features_resampled, labels_resampled = features_pca, labels
        print("未使用SMOTE，直接使用原始数据")
    
    return features_resampled, labels_resampled

# 4. 使用一种机器学习方法分析建模与模型评价，并对评价结果进行可视化
def model_and_evaluate(features, labels):
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(features, labels, test_size=0.2, random_state=42)
    # 建模
    model = LogisticRegression(max_iter=1000)  # 增加max_iter避免收敛警告
    model.fit(X_train, y_train)
    # 预测
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1]
    # 模型评价
    print(classification_report(y_test, y_pred))
    cm = confusion_matrix(y_test, y_pred)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.title('Confusion Matrix')
    plt.xlabel('Predicted Label')
    plt.ylabel('True Label')
    plt.savefig('confusion_matrix.png')
    plt.show()
    # ROC 曲线
    fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)
    roc_auc = auc(fpr, tpr)
    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, label=f'ROC curve (area = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.savefig('roc_curve.png')
    plt.show()
    return model

if __name__ == "__main__":
    # 获取数据
    data = get_stock_data()
    # 添加特征和标签
    data = add_features_and_labels(data)
    # 数据预处理
    features, labels = preprocess_data(data)
    # 建模和评价
    model = model_and_evaluate(features, labels)
    


# In[ ]:




